import cv2
import numpy as np


def order_points(pts):
    rect = np.zeros((4, 2), dtype="float32")
    s = pts.sum(axis=1)
    rect[0] = pts[np.argmin(s)]  # 左上
    rect[2] = pts[np.argmax(s)]  # 右下

    diff = np.diff(pts, axis=1)
    rect[1] = pts[np.argmin(diff)]  # 右上
    rect[3] = pts[np.argmax(diff)]  # 左下
    return rect


def get_id_card_corners(image_path):
    # 读取图像
    image = cv2.imread(image_path)
    if image is None:
        raise ValueError("图片读取失败，请检查路径是否正确")

    # 预处理
    gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
    blurred = cv2.GaussianBlur(gray, (5, 5), 0)
    edged = cv2.Canny(blurred, 50, 200)

    # 形态学闭操作
    kernel = cv2.getStructuringElement(cv2.MORPH_RECT, (5, 5))
    closed = cv2.morphologyEx(edged, cv2.MORPH_CLOSE, kernel)

    # 查找轮廓
    contours, _ = cv2.findContours(closed.copy(), cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
    if not contours:
        raise ValueError("未检测到任何轮廓")

    # 按面积降序排序
    contours = sorted(contours, key=cv2.contourArea, reverse=True)[:5]
    screenCnt = None

    for c in contours:
        peri = cv2.arcLength(c, True)
        approx = cv2.approxPolyDP(c, 0.02 * peri, True)
        if len(approx) == 4:
            screenCnt = approx
            break

    if screenCnt is None:
        raise ValueError("未找到四边形轮廓")

    # 获取并排序坐标点
    points = screenCnt.reshape(4, 2)
    ordered_points = order_points(points)

    return ordered_points.astype(int).tolist()


def perspective_transform(image, corners):
    """
    根据四个坐标点进行透视变换裁剪图像
    :param image: 原始图像（BGR格式）
    :param corners: 四个角点坐标 [[x1,y1], [x2,y2], [x3,y3], [x4,y4]]
    :return: 裁剪校正后的图像
    """
    # 将坐标转换为NumPy数组并指定精度
    src = np.array(corners, dtype=np.float32)

    # 解包有序坐标点
    (tl, tr, br, bl) = src  # 左上、右上、右下、左下

    # 计算新图像宽度（取上下边最大宽度）
    width_top = np.sqrt(((tr[0] - tl[0]) ** 2) + ((tr[1] - tl[1]) ** 2))
    width_bottom = np.sqrt(((br[0] - bl[0]) ** 2) + ((br[1] - bl[1]) ** 2))
    max_width = max(int(width_top), int(width_bottom))

    # 计算新图像高度（取左右边最大高度）
    height_right = np.sqrt(((br[0] - tr[0]) ** 2) + ((br[1] - tr[1]) ** 2))
    height_left = np.sqrt(((bl[0] - tl[0]) ** 2) + ((bl[1] - tl[1]) ** 2))
    max_height = max(int(height_right), int(height_left))

    # 定义目标四边形坐标（注意顺序一致）
    dst = np.array([
        [0, 0],
        [max_width - 1, 0],
        [max_width - 1, max_height - 1],
        [0, max_height - 1]
    ], dtype=np.float32)

    # 计算透视变换矩阵
    M = cv2.getPerspectiveTransform(src, dst)

    # 执行透视变换
    warped = cv2.warpPerspective(image, M, (max_width, max_height))

    return warped


# 示例使用
if __name__ == "__main__":
    try:
        # 读取原始图像
        input_path = "5.jpg"
        output_path = "cropped_id_card.jpg"
        image = cv2.imread(input_path)

        # 检测角点
        corners = get_id_card_corners(input_path)
        print("检测到坐标：", corners)

        # 执行裁切
        cropped = perspective_transform(image, corners)

        # 保存结果
        cv2.imwrite(output_path, cropped)
        print(f"裁切完成，结果已保存至 {output_path}")

        # 显示对比（可选）
        cv2.imshow("Original", image)
        cv2.imshow("Cropped", cropped)
        cv2.waitKey(0)
        cv2.destroyAllWindows()

    except Exception as e:
        print(f"处理出错：{str(e)}")